#include "bitstreamnalumodel.h"
#include <QDebug>
#include <QTime>
#include <QBrush>
#include <cstdio>
#include <pthread.h>

static QString nalu_type2str(int type) {
    switch (type) {
    case  NAL_UNIT_TYPE_UNSPECIFIED:
        return QString("UNSPEC");
    case  NAL_UNIT_TYPE_CODED_SLICE_NON_IDR:
        return QString("P");
    case  NAL_UNIT_TYPE_CODED_SLICE_DATA_PARTITION_A:
        return QString("DPA");
    case  NAL_UNIT_TYPE_CODED_SLICE_DATA_PARTITION_B:
        return QString("DPB");
    case  NAL_UNIT_TYPE_CODED_SLICE_DATA_PARTITION_C:
        return QString("DPC");
    case  NAL_UNIT_TYPE_CODED_SLICE_IDR:
        return QString("IDR");
    case  NAL_UNIT_TYPE_SEI:
        return QString("SEI");
    case  NAL_UNIT_TYPE_SPS:
        return QString("SPS");
    case  NAL_UNIT_TYPE_PPS:
        return QString("PPS");
    case  NAL_UNIT_TYPE_AUD:
        return QString("AUD");
    case  NAL_UNIT_TYPE_END_OF_SEQUENCE:
        return QString("END_OF_SEQUENCE");
    case  NAL_UNIT_TYPE_END_OF_STREAM:
        return QString("END_OF_STREAM");
    case  NAL_UNIT_TYPE_FILLER:
        return QString("FILLER");
    case  NAL_UNIT_TYPE_SPS_EXT:
        return QString("SPS_EXT");
    case  NAL_UNIT_TYPE_CODED_SLICE_AUX:
        return QString("SLICE_AUX");
    default:
        return QString("UNKNOWN");
    }

    return QString("UNKNOWN");
}

static Qt::GlobalColor nalu_type2color(int type) {
    switch (type) {
    case  NAL_UNIT_TYPE_UNSPECIFIED:
        return Qt::black;
    case  NAL_UNIT_TYPE_CODED_SLICE_NON_IDR:
        return Qt::white;
    case  NAL_UNIT_TYPE_CODED_SLICE_DATA_PARTITION_A:
        return Qt::white;
    case  NAL_UNIT_TYPE_CODED_SLICE_DATA_PARTITION_B:
        return Qt::white;
    case  NAL_UNIT_TYPE_CODED_SLICE_DATA_PARTITION_C:
        return Qt::white;
    case  NAL_UNIT_TYPE_CODED_SLICE_IDR:
        return Qt::red;
    case  NAL_UNIT_TYPE_SEI:
        return Qt::yellow;
    case  NAL_UNIT_TYPE_SPS:
        return Qt::yellow;
    case  NAL_UNIT_TYPE_PPS:
        return Qt::yellow;
    case  NAL_UNIT_TYPE_AUD:
        return Qt::yellow;
    case  NAL_UNIT_TYPE_END_OF_SEQUENCE:
        return Qt::yellow;
    case  NAL_UNIT_TYPE_END_OF_STREAM:
        return Qt::yellow;
    case  NAL_UNIT_TYPE_FILLER:
        return Qt::yellow;
    case  NAL_UNIT_TYPE_SPS_EXT:
        return Qt::yellow;
    case  NAL_UNIT_TYPE_CODED_SLICE_AUX:
        return Qt::yellow;
    default:
        return Qt::black;
    }

    return Qt::black;
}



BitStreamNALUModel::BitStreamNALUModel(QObject *parent)
     :QAbstractTableModel(parent), row_count(0) {
    /*
     *  h264_dbgfile is define in libh264bitstream.
     *  the function debug_nal write nalu info into this FILE* stream.
     */
    h264_dbgfile = open_memstream(&nal_info, &nal_info_size);
    if (!h264_dbgfile) {
        return;
    } else {
        h264_dbg_file = h264_dbgfile;
    }

    pthread_mutex_init(&parse_mutex, NULL);

    //parse_file("/opt/work/test.264");
}

BitStreamNALUModel::~BitStreamNALUModel() {
    if (h264_dbg_file) {
        fclose(h264_dbg_file);
        free(nal_info);
        nal_info_size = 0;
    }
    pthread_mutex_destroy(&parse_mutex);
}

int BitStreamNALUModel::rowCount(const QModelIndex & /*parent*/) const {
    return row_count;
}

int BitStreamNALUModel::columnCount(const QModelIndex & /*parent*/) const {
    return column_count;
}

QVariant BitStreamNALUModel::data(const QModelIndex &index, int role) const
{
    int row = index.row();
    int col = index.column();

    pthread_mutex_lock((pthread_mutex_t*)&parse_mutex);
    if(nalus.isEmpty()) {
        return QVariant();
    }

    if (row >= nalus.size()) {
        return QVariant();
    }

    struct naluInfo *nalu;
    nalu = nalus[row];
    if (!nalu) {
        return QVariant();
    }
    pthread_mutex_unlock((pthread_mutex_t*)&parse_mutex);

    switch(role){
    case Qt::DisplayRole:
        switch (col) {
        case 0:
            return QString("%1").arg(nalu->seq + 1);
            break;
        case 1:
            return QString("%1").arg(nalu->ref_idc);
            break;
        case 2:
            return nalu_type2str(nalu->type);
            break;
        case 3:
            return QString("%1").arg(nalu->len);
            break;
        default:
            return QVariant();
            break;
        }

        break;
    case Qt::BackgroundRole:

        Qt::GlobalColor color = nalu_type2color(nalu->type);
        QBrush redBackground(color);
        return redBackground;
    }
    return QVariant();
}

QVariant BitStreamNALUModel::headerData(int section, Qt::Orientation orientation,
                                        int role) const {
    if (role == Qt::DisplayRole) {
        if (orientation == Qt::Horizontal) {
            switch (section) {
            case 0:
                return QString("seq");
            case 1:
                return QString("nal_reference_idc");
            case 2:
                return QString("nal_type");
            case 3:
                return QString("nal_size");
            }
        }

#if 0
        if (orientation == Qt::Vertical) {
            switch (section) {
            default:
                return QString("%1").arg(section + 1);
            }
        }
#endif
    }
    return QVariant();
}

/*
 * This function is the cleanup routine when the parse thread is canceled.
 */
void BitStreamNALUModel::_thread_cleanup(void *param) {

    class BitStreamNALUModel *naluModel;
    naluModel = (class BitStreamNALUModel*)param;

    //free the stream handle.
    h264_free(naluModel->h264_stream);
    //close the FILE* handle for the file being parsed
    fclose(naluModel->infile);
    //reset the NALU info buffer
    rewind(naluModel->h264_dbg_file);
    //remove all the content in the view.
    naluModel->removeRows(0, naluModel->rowCount());
    // free all the memory alloced to store the nalus info.
    for (auto it = naluModel->nalus.begin(); it != naluModel->nalus.end(); ++it) {
        delete *it;
    }
    // clear the vector, and release its memory
    naluModel->nalus.clear();
    QVector<struct naluInfo*>().swap(naluModel->nalus);

    return;
}

void *BitStreamNALUModel::_thread_parser_file(void *param) {

    int seq;
    class BitStreamNALUModel *naluModel;
    QByteArray ba;

    pthread_setcanceltype(PTHREAD_CANCEL_DEFERRED, NULL);

    seq = 0;
    naluModel = (class BitStreamNALUModel*)param;

    ba = naluModel->filename.toLatin1();
    naluModel->infile = fopen(ba.data(), "r");
    if (!naluModel->infile) {
        return NULL;
    }

    naluModel->h264_stream = h264_new();
    if (!naluModel->h264_stream){
        qDebug() << "Failed to create a new h264 stream";
        return NULL;
    }

    if(naluModel->nalus.size()) {
        for (auto it = naluModel->nalus.begin(); it != naluModel->nalus.end(); ++it) {
            delete *it;
        }
        naluModel->nalus.clear();
        QVector<struct naluInfo*>().swap(naluModel->nalus);
    }

    size_t rsz = 0;
    size_t sz = 0;
    int64_t off = 0;
    uint8_t buf[BUFSIZ];
    uint8_t* p = buf;

    int nal_start, nal_end;
    int row;
    naluModel->removeRows(0, naluModel->rowCount());
    int last_row_index = naluModel->rowCount();

#define INSERT_ROWS \
    do { \
            qDebug() << "last_row_index: " << last_row_index; \
            QModelIndex topLeft = naluModel->createIndex(last_row_index, 0); \
            QModelIndex bottomRight = naluModel->createIndex(row, 2); \
            qDebug() << "row: " << row; \
            qDebug() << "deta: " << row - last_row_index; \
            naluModel->insertRows(last_row_index, row - last_row_index);\
            emit naluModel->dataChanged(topLeft, bottomRight); \
            last_row_index = naluModel->rowCount(); \
        } while (0)


    while (1)
    {
        rsz = fread(buf + sz, 1, BUFSIZ - sz, naluModel->infile);
        if (rsz == 0)
        {
            if (ferror(naluModel->infile)) {
                fprintf( stderr, "!! Error: read failed: %s \n", strerror(errno));
                break;
            }
            break;  // if (feof(infile))
        }

        sz += rsz;

        while (find_nal_unit(p, sz, &nal_start, &nal_end) > 0)
        {
            p += nal_start;
            read_nal_unit(naluModel->h264_stream, p, nal_end - nal_start);

            debug_nal(naluModel->h264_stream, naluModel->h264_stream->nal);
            struct naluInfo *nalu;
            nalu = new struct naluInfo;
            row = seq++;
            nalu->seq = row;
            nalu->ref_idc = naluModel->h264_stream->nal->nal_ref_idc;
            nalu->type = naluModel->h264_stream->nal->nal_unit_type;
            nalu->len = nal_end - nal_start;
            fflush(h264_dbgfile);
            nalu->info = QString(naluModel->nal_info);
            pthread_setcancelstate(PTHREAD_CANCEL_DISABLE, NULL);
            pthread_mutex_lock(&naluModel->parse_mutex);
            qDebug() << nalu->seq << endl;
            naluModel->nalus.push_back(nalu);
            pthread_mutex_unlock(&naluModel->parse_mutex);

            pthread_setcancelstate(PTHREAD_CANCEL_ENABLE, NULL);
            pthread_cleanup_push(BitStreamNALUModel::_thread_cleanup, naluModel);
            pthread_testcancel();
            pthread_cleanup_pop(0);

            if (seq % 500 == 0) {
                INSERT_ROWS;
            }

            rewind(h264_dbgfile);
            p += (nal_end - nal_start);
            sz -= nal_end;
        }

        // if no NALs found in buffer, discard it
        if (p == buf)
        {
            p = buf + sz;
            sz = 0;
        }

        memmove(buf, p, sz);
        off += p - buf;
        p = buf;
    }

    INSERT_ROWS;

    h264_free(naluModel->h264_stream);
    fclose(naluModel->infile);
    return NULL;
}

void BitStreamNALUModel::parse_file(QString fname) {
    /*
     * The main part of this function is from libh264bitstream, file h264_analyze.c
     */

    if (filename == fname) {
        return;
    } else {
        if(!filename.isEmpty()) {
            pthread_cancel(parse_tid);
            pthread_join(parse_tid, NULL);
        }
        filename = fname;
    }

    pthread_create(&parse_tid, NULL, BitStreamNALUModel::_thread_parser_file, this);
}

struct naluInfo *BitStreamNALUModel::naluByIndex(int row)
{
    if (row >= nalus.size()) {
        return Q_NULLPTR;
    }

    return nalus[row];
}

bool BitStreamNALUModel::insertRows(int row, int count,
                                    const QModelIndex &parent) {
    qDebug() << "insert row: " << row << ", count: " << count;
    beginInsertRows(parent, row, row + count - 1);
    row_count += count;
    endInsertRows();
    return true;
}
bool BitStreamNALUModel::removeRows(int row, int count, const QModelIndex &parent ) {
    qDebug() << "remove row: " << row << ", count: " << count;
    beginRemoveRows(parent, row, row + count - 1);
    row_count -= count;
    endRemoveRows();
    return true;
}

void BitStreamNALUModel::resize_rows()
{
    return;
}

void BitStreamNALUModel::resize_columns()
{
    return;
}

QModelIndex BitStreamNALUModel::get_index(int row, int c)
{
    return createIndex(row, c);
}
